from flask import Blueprint, render_template, redirect, url_for, flash
from python.model.user import User
from wtforms import StringField, PasswordField, SubmitField, SelectField  # 添加 SelectField
from wtforms.validators import DataRequired, Email, EqualTo
from python.extensions import db, login_manager
from flask_wtf import FlaskForm
from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user
login_manager.login_view = 'auth.login'

auth_bp = Blueprint('auth', __name__, url_prefix='/auth')

# 登录表单
class LoginForm(FlaskForm):
    username = StringField('用户名', validators=[DataRequired()])
    password = PasswordField('密码', validators=[DataRequired()])
    # 添加登录类型选择字段
    user_type = SelectField('登录类型', choices=[
        ('user', '普通用户'),
        ('admin', '管理员')
    ], validators=[DataRequired()])
    submit = SubmitField('登录')

# 注册表单
class RegisterForm(FlaskForm):
    username = StringField('用户名', validators=[DataRequired()])
    email = StringField('邮箱', validators=[DataRequired(), Email()])
    password = PasswordField('密码', validators=[DataRequired()])
    password2 = PasswordField('确认密码',
                              validators=[DataRequired(), EqualTo('password')])
    admin_code = StringField('管理员验证码')
    submit = SubmitField('注册')

@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))



@auth_bp.route('/login', methods=['GET', 'POST'])
def login():
    form = LoginForm()
    if form.validate_on_submit():
        user = User.query.filter_by(username=form.username.data).first()

        # 检查用户是否存在且密码正确
        if user and user.check_password(form.password.data):
            # 处理管理员登录
            if form.user_type.data == 'admin':
                if user.is_admin:
                    login_user(user)
                    return redirect(url_for('main.index'))  # 统一重定向到 main.index
            # 处理普通用户登录
            elif form.user_type.data == 'user':
                login_user(user)
                return redirect(url_for('main.index'))  # 重定向到 main.index

        # 用户名或密码错误
        flash('用户名或密码错误')

    return render_template('auth/login.html', form=form)


from sqlalchemy.exc import IntegrityError

@auth_bp.route('/register', methods=['GET', 'POST'])
def register():
    form = RegisterForm()
    if form.validate_on_submit():
        existing_user = User.query.filter(
            (User.username == form.username.data) | (User.email == form.email.data)
        ).first()
        if existing_user:
            flash('用户名或者邮箱已存在')
            return redirect(url_for('auth.register'))

        new_user = User(
            username=form.username.data,
            email=form.email.data
        )
        new_user.set_password(form.password.data)

        # 验证管理员验证码并设置管理员权限
        if form.admin_code.data == "2216030305":  # 管理员验证码
            new_user.is_admin = True
        else:
            new_user.is_admin = False  # 默认为普通用户

        db.session.add(new_user)
        db.session.commit()

        flash('注册成功! 请登录')
        return redirect(url_for('auth.login'))  # 注册成功后重定向到登录页面
    return render_template('auth/register.html', form=form)

@auth_bp.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('auth.login'))